#include "postprocess/matcher/matcher.h"

// #include <float.h>

#include <queue>

float Mathcer::s_match_2DIOU_minimum_ = 0.1f;
float Mathcer::s_match_distance_maximum_ = 1.8f;

Mathcer::Mathcer(const ProcessorClass &flag) : processor_flag_(flag) { init(); }

Mathcer::Mathcer(const ProcessorClass &flag,
                 const PostProcessorSettings &settings)
    : processor_flag_(flag), settings_(settings) {
  init();
}

Mathcer::~Mathcer() {}

void Mathcer::init() {
  s_match_2DIOU_minimum_ = 0.1f;
  s_match_distance_maximum_ = 1.8f;

  // 载入cost计算与比较函数对象
  // 0 -- 汽车类
  // 1 -- 行人类
  cost_esti_func_map_.insert(
      {0, std::make_pair(&Cost::ComputerIOU2DWithPos, &Mathcer::greater_than)});
  cost_esti_func_map_.insert(
      {1, std::make_pair(&Cost::ComputerCartesian2DDistanceWithPos,
                         &Mathcer::less_than)});
}

void Mathcer::Match(std::vector<TrackedObjectPtr> *objects,
                    const std::vector<ObjectTrackPtr> &tracks,
                    const std::vector<Eigen::VectorXf> &tracks_predict,
                    std::vector<std::pair<int, int>> *assignments,
                    std::vector<int> *unassigned_tracks,
                    std::vector<int> *unassigned_objects) {
  // A. computing association matrix
  Eigen::MatrixXf association_mat(tracks.size(), objects->size());
  compute_associate_matrix(tracks, tracks_predict, (*objects),
                           &association_mat);

  // B. computing connected components
  float threshold = 0.0f;
  std::vector<std::vector<std::int32_t>> object_components;  // object联通id
  std::vector<std::vector<std::int32_t>> track_components;   // track联通id
  if (ProcessorClass::Vehicle == processor_flag_) {          // 计算汽车
    threshold = settings_.vehicle_2diou_threshold;
  } else if (ProcessorClass::Pedestrian == processor_flag_) {  // 计算行人
    threshold = settings_.pedestrian_distance_threshold;
  }
  compute_connected_components(association_mat, threshold, &track_components,
                               &object_components);

  // C. matching each sub-graph
  assignments->clear();
  unassigned_tracks->clear();
  unassigned_objects->clear();
  for (size_t i = 0; i < track_components.size(); i++) {
    std::vector<std::pair<int, int>> sub_assignments;
    std::vector<int> sub_unassigned_tracks;
    std::vector<int> sub_unassigned_objects;
    match_in_components(association_mat, track_components[i],
                        object_components[i], threshold, &sub_assignments,
                        &sub_unassigned_tracks, &sub_unassigned_objects);
    for (size_t j = 0; j < sub_assignments.size(); ++j) {
      int track_id = sub_assignments[j].first;
      int object_id = sub_assignments[j].second;
      assignments->push_back(sub_assignments[j]);
      float association_score = association_mat(track_id, object_id);
      (*objects)[object_id]->association_score = association_score;
    }
    for (size_t j = 0; j < sub_unassigned_tracks.size(); ++j) {
      unassigned_tracks->push_back(sub_unassigned_tracks[j]);
    }
    for (size_t j = 0; j < sub_unassigned_objects.size(); ++j) {
      unassigned_objects->push_back(sub_unassigned_objects[j]);
    }
  }
}

void Mathcer::Match(std::vector<TrackedObjectPtr> *objects,
                    const std::vector<ObjectTrackPtr> &tracks,
                    const std::vector<Eigen::VectorXf> &tracks_predict,
                    std::vector<std::pair<int, int>> *assignments,
                    std::vector<int> *unassigned_tracks,
                    std::vector<int> *unassigned_objects,
                    const Eigen::Matrix4f &pos) {
  // A. computing association matrix
  Eigen::MatrixXf association_mat(tracks.size(), objects->size());
  compute_associate_matrix(tracks, tracks_predict, (*objects), &association_mat,
                           pos);

  // std::cout
  //     << "-------------------------match start-----------------------------"
  //     << std::endl;
  // for (auto iter = objects->begin(); iter != objects->end(); ++iter) {
  //   std::cout << "---------new objects------------" << std::endl;
  //   std::cout << "label = " << (*iter)->label << std::endl;
  //   std::cout << "center = " << (*iter)->center.x() << " "
  //             << (*iter)->center.y() << " " << (*iter)->center.z() <<
  //             std::endl;

  //   std::cout << "vel = " << (*iter)->velocity.x() << " "
  //             << (*iter)->velocity.y() << " " << (*iter)->velocity.z()
  //             << std::endl;
  //   std::cout << "size = " << (*iter)->size.x() << " " << (*iter)->size.y()
  //             << " " << (*iter)->size.z() << std::endl;
  // }

  // for (auto iter = tracks.begin(); iter != tracks.end(); ++iter) {
  //   std::cout << "---------tracks--------------" << std::endl;
  //   std::cout << "label = " << (*iter)->current_object_->label << std::endl;
  //   std::cout << "center = " << (*iter)->current_object_->center.x() << " "
  //             << (*iter)->current_object_->center.y() << " "
  //             << (*iter)->current_object_->center.z() << std::endl;
  //   std::cout << "size = " << (*iter)->current_object_->size.x() << " "
  //             << (*iter)->current_object_->size.y() << " "
  //             << (*iter)->current_object_->size.z() << std::endl;
  // }

  // for (auto iter = tracks_predict.begin(); iter != tracks_predict.end();
  //      ++iter) {
  //   std::cout << "---------tracks_predict--------------" << std::endl;
  //   std::cout << "predict = " << (*iter)[0] << " " << (*iter)[1] << " "
  //             << (*iter)[2] << " " << (*iter)[3] << " " << (*iter)[4] << " "
  //             << (*iter)[5] << " " << (*iter)[6] << std::endl;
  // }

  // std::cout << "---------association_mat" << std::endl;
  // std::cout << "association_mat = " << association_mat << std::endl;

  // std::cout << "-----------------------match end-- -- -- -- -- -- -- -- --
  // --"
  //              "-- -- -- -- -"
  //           << std::endl;

  // B. computing connected components
  float threshold = 0.0f;
  std::vector<std::vector<std::int32_t>> object_components;  // object联通id
  std::vector<std::vector<std::int32_t>> track_components;   // track联通id
  if (ProcessorClass::Vehicle == processor_flag_) {          // 计算汽车
    threshold = s_match_2DIOU_minimum_;
    // compute_connected_components(association_mat, s_match_2DIOU_minimum_,
    //                              &track_components, &object_components);
  } else if (ProcessorClass::Pedestrian == processor_flag_) {  // 计算行人
    threshold = s_match_distance_maximum_;
    // compute_connected_components(association_mat, s_match_distance_maximum_,
    //                              &track_components, &object_components);
  }
  compute_connected_components(association_mat, threshold, &track_components,
                               &object_components);

  // C. matching each sub-graph
  assignments->clear();
  unassigned_tracks->clear();
  unassigned_objects->clear();
  for (size_t i = 0; i < track_components.size(); i++) {
    std::vector<std::pair<int, int>> sub_assignments;
    std::vector<int> sub_unassigned_tracks;
    std::vector<int> sub_unassigned_objects;
    match_in_components(association_mat, track_components[i],
                        object_components[i], threshold, &sub_assignments,
                        &sub_unassigned_tracks, &sub_unassigned_objects);
    for (size_t j = 0; j < sub_assignments.size(); ++j) {
      int track_id = sub_assignments[j].first;
      int object_id = sub_assignments[j].second;
      assignments->push_back(sub_assignments[j]);
      float association_score = association_mat(track_id, object_id);
      (*objects)[object_id]->association_score = association_score;
    }
    for (size_t j = 0; j < sub_unassigned_tracks.size(); ++j) {
      unassigned_tracks->push_back(sub_unassigned_tracks[j]);
    }
    for (size_t j = 0; j < sub_unassigned_objects.size(); ++j) {
      unassigned_objects->push_back(sub_unassigned_objects[j]);
    }
  }
}

void Mathcer::compute_associate_matrix(
    std::vector<ObjectTrackPtr> const &tracks,
    std::vector<Eigen::VectorXf> const &tracks_predict,
    std::vector<TrackedObjectPtr> const &new_objects,
    Eigen::MatrixXf *association_mat) {
  for (size_t i = 0; i < tracks.size(); i++) {
    for (size_t j = 0; j < new_objects.size(); j++) {
      (*association_mat)(i, j) = compute_connection_evaluation(
          tracks[i], tracks_predict[i], new_objects[j]);
    }
  }
}

void Mathcer::compute_associate_matrix(
    std::vector<ObjectTrackPtr> const &tracks,
    std::vector<Eigen::VectorXf> const &tracks_predict,
    std::vector<TrackedObjectPtr> const &new_objects,
    Eigen::MatrixXf *association_mat, const Eigen::Matrix4f &pos) {
  for (size_t i = 0; i < tracks.size(); i++) {
    for (size_t j = 0; j < new_objects.size(); j++) {
      (*association_mat)(i, j) = compute_connection_evaluation(
          tracks[i], tracks_predict[i], new_objects[j], pos);
    }
  }
}

float Mathcer::compute_connection_evaluation(
    ObjectTrackPtr const &track, Eigen::VectorXf const &track_predict,
    TrackedObjectPtr const &new_object) {
  float result = 0.0f;
  // vehicle:2D IOU    pedestrain:location
  // result = cost_esti_func_map_[static_cast<int>(processor_flag_)].first(
  //     track, track_predict, new_object);
  return result;
}

float Mathcer::compute_connection_evaluation(
    ObjectTrackPtr const &track, Eigen::VectorXf const &track_predict,
    TrackedObjectPtr const &new_object, const Eigen::Matrix4f &pos) {
  float result = 0.0f;
  // vehicle:2D IOU    pedestrain:location
  result = cost_esti_func_map_[static_cast<int>(processor_flag_)].first(
      track, track_predict, new_object, pos);
  return result;
}

float Mathcer::VectorCosTheta2fXy(Eigen::Vector3f const &v1,
                                  Eigen::Vector3f const &v2) {
  float v1_len = sqrt(v1.x() * v1.x() + v1.y() * v1.y());
  float v2_len = sqrt(v2.x() * v2.x() + v2.y() * v2.y());
  float cos_theta{0.0f};
  if (v1_len > FLT_EPSILON && v2_len > FLT_EPSILON) {
    cos_theta = (v1.x() * v2.x() + v1.y() * v2.y()) / (v1_len * v2_len);
  }
  return cos_theta;
}

void Mathcer::compute_connected_components(
    const Eigen::MatrixXf &association_mat, const float connected_threshold,
    std::vector<std::vector<int32_t>> *track_components,
    std::vector<std::vector<int32_t>> *object_components) {
  // Compute connected components within given threshold
  int no_track = association_mat.rows();
  int no_object = association_mat.cols();
  std::vector<std::vector<int>> nb_graph;
  nb_graph.resize(no_track + no_object);
  for (int i = 0; i < no_track; i++) {
    for (int j = 0; j < no_object; j++) {
      //      if (association_mat(i, j) > connected_threshold) {
      if (cost_esti_func_map_[static_cast<int>(processor_flag_)].second(
              association_mat(i, j), connected_threshold)) {
        nb_graph[i].push_back(no_track + j);
        nb_graph[j + no_track].push_back(i);
      }
    }
  }

  std::vector<std::vector<int>> components;
  connected_component_analysis(nb_graph, &components);
  track_components->clear();
  track_components->resize(components.size());
  object_components->clear();
  object_components->resize(components.size());
  for (size_t i = 0; i < components.size(); i++) {
    for (size_t j = 0; j < components[i].size(); j++) {
      int id = components[i][j];
      if (id < no_track) {
        (*track_components)[i].push_back(id);
      } else {
        id -= no_track;
        (*object_components)[i].push_back(id);
      }
    }
  }
}

void Mathcer::connected_component_analysis(
    const std::vector<std::vector<int>> &graph,
    std::vector<std::vector<int>> *components) {
  int no_item = graph.size();
  std::vector<int> visited;
  visited.resize(no_item, 0);
  std::queue<int> que;
  std::vector<int> component;
  components->clear();

  for (int i = 0; i < no_item; ++i) {
    if (visited[i]) {
      continue;
    }
    component.push_back(i);
    que.push(i);
    visited[i] = 1;
    while (!que.empty()) {
      int id = que.front();  // 访问第一个元素
      que.pop();             // 删除第一个元素
      for (size_t j = 0; j < graph[id].size(); ++j) {
        int nb_id = graph[id][j];
        if (visited[nb_id] == 0) {
          component.push_back(nb_id);
          que.push(nb_id);
          visited[nb_id] = 1;
        }
      }
    }
    components->push_back(component);
    component.clear();
  }
}

void Mathcer::match_in_components(
    const Eigen::MatrixXf &association_mat,
    const std::vector<int32_t> &track_component,
    const std::vector<int32_t> &object_component,
    const float &connection_threshold,
    std::vector<TrackObjectPair> *sub_assignments,
    std::vector<int32_t> *sub_unassigned_tracks,
    std::vector<int32_t> *sub_unassigned_objects) {
  sub_assignments->clear();
  sub_unassigned_tracks->clear();
  sub_unassigned_objects->clear();
  // A. failted to match if either components is empty
  if (track_component.empty()) {
    for (size_t i = 0; i < object_component.size(); ++i) {
      sub_unassigned_objects->push_back(object_component[i]);
    }
  }
  if (object_component.empty()) {
    for (size_t i = 0; i < track_component.size(); ++i) {
      sub_unassigned_tracks->push_back(track_component[i]);
    }
  }
  if (track_component.empty() || object_component.empty()) return;
  // B. if components perfectly match
  if (track_component.size() == 1 && object_component.size() == 1) {
    int track_id = track_component[0];
    int object_id = object_component[0];

    if (cost_esti_func_map_[static_cast<int>(processor_flag_)].second(
            association_mat(track_id, object_id), connection_threshold)) {
      sub_assignments->push_back(std::make_pair(track_id, object_id));
    } else {
      sub_unassigned_objects->push_back(object_id);
      sub_unassigned_tracks->push_back(track_id);
    }
    return;
  }
  // C. multi object track match
  std::vector<int> track_local2global;
  std::vector<int> object_local2global;
  std::vector<std::pair<int, int>> local_assignments;
  std::vector<int> local_unassigned_tracks;
  std::vector<int> local_unassigned_objects;
  Eigen::MatrixXf local_association_mat(track_component.size(),
                                        object_component.size());
  track_local2global.resize(track_component.size());
  object_local2global.resize(object_component.size());
  for (size_t i = 0; i < track_component.size(); ++i) {
    track_local2global[i] = track_component[i];
  }
  for (size_t i = 0; i < object_component.size(); ++i) {
    object_local2global[i] = object_component[i];
  }
  for (size_t i = 0; i < track_component.size(); ++i) {
    for (size_t j = 0; j < object_component.size(); ++j) {
      int track_id = track_component[i];
      int object_id = object_component[j];
      local_association_mat(i, j) = association_mat(track_id, object_id);
    }
  }
  local_assignments.resize(local_association_mat.cols());
  local_unassigned_tracks.assign(local_association_mat.rows(), -1);
  local_unassigned_objects.assign(local_association_mat.cols(), -1);
  //  assign_objects_to_tracks(local_association_mat, s_match_distance_maximum_,
  //                           &local_assignments, &local_unassigned_tracks,
  //                           &local_unassigned_objects);
  assign_objects_to_tracks(local_association_mat, connection_threshold,
                           &local_assignments, &local_unassigned_tracks,
                           &local_unassigned_objects);
  for (size_t i = 0; i < local_assignments.size(); ++i) {
    int global_track_id = track_local2global[local_assignments[i].first];
    int global_object_id = object_local2global[local_assignments[i].second];
    sub_assignments->push_back(
        std::make_pair(global_track_id, global_object_id));
  }
  for (size_t i = 0; i < local_unassigned_tracks.size(); ++i) {
    int global_track_id = track_local2global[local_unassigned_tracks[i]];
    sub_unassigned_tracks->push_back(global_track_id);
  }
  for (size_t i = 0; i < local_unassigned_objects.size(); ++i) {
    int global_object_id = object_local2global[local_unassigned_objects[i]];
    sub_unassigned_objects->push_back(global_object_id);
  }
}

void Mathcer::assign_objects_to_tracks(
    Eigen::MatrixXf const &association_mat,
    const double assign_distance_maximum,
    std::vector<TrackObjectPair> *assignments,
    std::vector<std::int32_t> *unassigned_tracks,
    std::vector<std::int32_t> *unassigned_objects) noexcept {
  std::vector<std::int32_t> tracks_idx;
  std::vector<std::int32_t> objects_idx;
  std::int32_t no_track = association_mat.rows();
  std::int32_t no_object = association_mat.cols();
  // A. build cost function
  std::vector<std::vector<double>> cost(no_track + no_object);
  for (std::int32_t i = 0; i < no_track; ++i) {
    cost[i].resize(no_object);
    for (std::int32_t j = 0; j < no_object; ++j) {
      cost[i][j] = association_mat(i, j);
    }
  }
  for (std::int32_t i = 0; i < no_object; ++i) {
    cost[i + no_track].resize(no_object);
    for (std::int32_t j = 0; j < no_object; j++) {
      if (j == i) {
        cost[i + no_track][j] = assign_distance_maximum * 1.2f;
      } else {
        cost[i + no_track][j] = 999999.0f;
      }
    }
  }

  // B. HungarianOptimizer
  HungarianOptimizer hungarian_optimizer(cost);
  hungarian_optimizer.minimize(&tracks_idx, &objects_idx);

  // C. identify assignments, unassigned_tracks, unassigned_objects
  std::int32_t assignments_num = 0;
  std::vector<bool> tracks_used(no_track + no_object, false);
  std::vector<bool> objects_used(no_object, false);

  for (size_t i = 0; i < tracks_idx.size(); ++i) {
    if (tracks_idx[i] < 0 || tracks_idx[i] >= no_track || objects_idx[i] < 0 ||
        objects_idx[i] >= no_object) {
      continue;
    }
    //    if (association_mat(tracks_idx[i], objects_idx[i]) <
    //        assign_distance_maximum) {
    if (!cost_esti_func_map_[static_cast<int>(processor_flag_)].second(
            association_mat(tracks_idx[i], objects_idx[i]),
            assign_distance_maximum)) {
      (*assignments)[assignments_num++] =
          std::make_pair(tracks_idx[i], objects_idx[i]);
      tracks_used[tracks_idx[i]] = true;
      objects_used[objects_idx[i]] = true;
    }
  }
  assignments->resize(assignments_num);

  unassigned_tracks->resize(association_mat.rows());
  std::int32_t unassigned_tracks_num = 0;
  for (std::int32_t i = 0; i < association_mat.rows(); ++i) {
    if (tracks_used[i] == false) {
      (*unassigned_tracks)[unassigned_tracks_num++] = i;
    }
  }
  unassigned_tracks->resize(unassigned_tracks_num);

  unassigned_objects->resize(association_mat.cols());
  std::int32_t unassigned_objects_num = 0;
  for (std::int32_t i = 0; i < association_mat.cols(); ++i) {
    if (objects_used[i] == false) {
      (*unassigned_objects)[unassigned_objects_num++] = i;
    }
  }
  unassigned_objects->resize(unassigned_objects_num);
}
